import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
matplotlib.use('Agg')

# 定义色卡
dbz_list = ["#0000F6", "#01A0F6", "#00ECEC", "#01FF00", "#00C800", "#019000",
            "#FFFF00", "#E7C000", "#FF9000", "#FF0000", "#D60000", "#C00000",
            "#FF00F0", "#780084", "#AD90F0"]
my_cmap = LinearSegmentedColormap.from_list('dbz', dbz_list, N=15)


def plot_single_heatmap(img, fname, format="png", transparent=True):
    """

    @param img: 输入的图像数据  numpy  都是在0-70范围内
    @param fname: 保存的名称
    @param format: 保存格式  png\svg
    @param black_background:   是否设置背景板是黑色
    @param transparent:
    @return:
    """
    mask = img <= 5  #设置掩码的区域为小于5
    img_mask = np.ma.array(img, mask=mask)
    fig = plt.figure(
        figsize=(1, 1),  # 画布大小， (Width, height)，单位：inches
        dpi=400,
        # 图像分辨率
        facecolor='w',  # 背景色，是画布底下图层的颜色，默认白色
        linewidth=0,  # 画布边框的宽度，默认 0，无边框
        edgecolor='w',  # 边框颜色，要设置 linewidth 较大时才能看见，默认白色
        frameon=False,  # 默认值 True 为绘制边框，如果为 False 则不绘制边框
    )

    axes = fig.add_subplot(111)  # 添加子图，这里 111 表示就一张图，返回该子图的坐标系对象

    axes.contourf(img_mask, origin='upper', levels=np.arange(0, 76, 5), cmap=my_cmap)
    # 关闭坐标轴
    axes.set_axis_off()
    plt.subplots_adjust(left=0, bottom=0, right=1, top=1, hspace=0.1, wspace=0.1)
    # 保存图片
    plt.savefig(fname, dpi=400, format=format, transparent=transparent)
    plt.close()
